// 🚀 零拷贝流式文件传输服务器 - Go实现
// 支持3MB分片接收 + MD5校验 + 自动合并 + 零拷贝优化
package main

import (
	"bufio"
	"crypto/md5"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
	"math/rand"
)

// 📦 常量定义
const (
	PORT              = ":8888"  // 修改端口为8888
	CHUNK_SIZE        = 3 * 1024 * 1024  // 3MB分片大小
	TEMP_DIR          = "temp/"          // 临时文件目录
	BUFFER_SIZE       = 64 * 1024        // 64KB缓冲区
)

// 📋 传输任务结构体
type TransferTask struct {
	TaskID         string
	FileName       string
	FileSize       int64
	TotalChunks    int
	MD5Hash        string
	ReceivedChunks map[int]*os.File  // 已接收的分片文件
	mu             sync.RWMutex      // 读写锁保护
	ReceivedCount  int               // 已接收分片计数
	LastProgress   int               // 上次显示的进度百分比
}

// 🌐 全局变量
var (
	transferTasks = make(map[string]*TransferTask)  // 传输任务映射
	tasksMutex    = sync.RWMutex{}                  // 任务映射锁
	connectionCount = 0                             // 连接计数器
	connectionMutex = sync.Mutex{}                  // 连接计数器锁
)

// 🚀 主函数入口
func main() {
	// 📁 创建临时目录
	if err := os.MkdirAll(TEMP_DIR, 0755); err != nil {
		fmt.Printf("❌ 创建临时目录失败: %v\n", err)
		return
	}

	// 📡 启动TCP服务器
	listener, err := net.Listen("tcp", PORT)
	if err != nil {
		fmt.Printf("❌ 服务器启动失败: %v\n", err)
		return
	}
	defer listener.Close()

	fmt.Println("🚀 启动零拷贝流式文件传输服务器")
	fmt.Printf("🚀 零拷贝流式文件传输服务器启动!\n")
	fmt.Printf("🌍 监听端口: %s\n", PORT)
	fmt.Printf("📁 保存目录: 当前目录\n")
	fmt.Printf("💾 内存优化: 流式写入，零数据缓存\n")
	fmt.Printf("🔌 等待客户端连接...\n")

	// 🔁 循环接受客户端连接
	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Printf("❌ 接受连接失败: %v\n", err)
			continue
		}

		// 🔢 增加连接计数
		connectionMutex.Lock()
		connectionCount++
		connID := connectionCount
		connectionMutex.Unlock()

		// 🧵 为每个连接启动独立协程处理
		go handleConnection(conn, connID)
	}
}

// 🧠 处理客户端连接
func handleConnection(conn net.Conn, connID int) {
	defer conn.Close()
	
	// 📡 获取客户端地址
	clientAddr := conn.RemoteAddr().String()

	// 📖 创建带缓冲的读取器
	reader := bufio.NewReader(conn)

	// 📬 循环读取消息
	for {
		message, err := reader.ReadString('\n')
		if err != nil {
			if err != io.EOF {
				fmt.Printf("❌ [%s] 读取消息失败: %v\n", clientAddr, err)
			}
			break
		}

		// 🔪 解析消息类型
		parts := strings.Split(strings.TrimSpace(message), "|")
		if len(parts) == 0 {
			continue
		}

		switch parts[0] {
		case "HANDSHAKE":
			handleHandshake(conn, reader, clientAddr, parts)
		case "CHUNK":
			handleChunk(conn, reader, clientAddr, parts)
		case "TRANSFER_COMPLETE":
			handleTransferComplete(conn, clientAddr, parts)
		}
	}
}

// 🤝 处理握手消息
func handleHandshake(conn net.Conn, reader *bufio.Reader, clientAddr string, parts []string) {
	if len(parts) != 6 {
		fmt.Printf("❌ [%s] 握手消息格式错误\n", clientAddr)
		return
	}

	// 📄 解析握手信息
	taskID := parts[1]
	fileName := parts[2]
	fileSize, err := strconv.ParseInt(parts[3], 10, 64)
	if err != nil {
		fmt.Printf("❌ [%s] 文件大小解析失败: %v\n", clientAddr, err)
		return
	}
	
	totalChunks, err := strconv.Atoi(parts[4])
	if err != nil {
		fmt.Printf("❌ [%s] 分片总数解析失败: %v\n", clientAddr, err)
		return
	}
	
	md5Hash := parts[5]

	// 📋 创建传输任务
	task := &TransferTask{
		TaskID:         taskID,
		FileName:       fileName,
		FileSize:       fileSize,
		TotalChunks:    totalChunks,
		MD5Hash:        md5Hash,
		ReceivedChunks: make(map[int]*os.File),
		ReceivedCount:  0,
		LastProgress:   0,
	}

	// 🔐 注册传输任务
	tasksMutex.Lock()
	transferTasks[taskID] = task
	tasksMutex.Unlock()

	fmt.Printf("📁 [%s] 文件: %s (%.2f MB)\n", clientAddr, fileName, float64(fileSize)/(1024*1024))
	fmt.Printf("📦 [%s] 分片总数: %d, MD5: %s\n", clientAddr, totalChunks, md5Hash)
	fmt.Printf("💾 [%s] 零拷贝模式: 分片直接写入临时文件\n", clientAddr)

	// ✅ 发送握手确认
	_, err = conn.Write([]byte("HANDSHAKE_OK\n"))
	if err != nil {
		fmt.Printf("❌ [%s] 发送握手确认失败: %v\n", clientAddr, err)
		return
	}

	fmt.Printf("✅ [%s] 控制连接建立成功，等待分片传输\n", clientAddr)
}

// 📦 处理分片消息
func handleChunk(conn net.Conn, reader *bufio.Reader, clientAddr string, parts []string) {
	if len(parts) != 4 {
		fmt.Printf("❌ [%s] 分片消息格式错误\n", clientAddr)
		return
	}

	// 🔢 解析分片信息
	taskID := parts[1]
	chunkIndex, err := strconv.Atoi(parts[2])
	if err != nil {
		fmt.Printf("❌ [%s] 分片索引解析失败: %v\n", clientAddr, err)
		return
	}
	
	chunkSize, err := strconv.Atoi(parts[3])
	if err != nil {
		fmt.Printf("❌ [%s] 分片大小解析失败: %v\n", clientAddr, err)
		return
	}

	// 🔍 查找对应的传输任务
	tasksMutex.RLock()
	task, exists := transferTasks[taskID]
	tasksMutex.RUnlock()
	
	if !exists {
		fmt.Printf("❌ [%s] 未找到对应的传输任务\n", clientAddr)
		return
	}

	// 📁 创建临时文件目录（使用任务ID确保唯一性）
	tempDir := filepath.Join(TEMP_DIR, fmt.Sprintf("%s_%s", task.FileName, taskID))
	if err := os.MkdirAll(tempDir, 0755); err != nil {
		fmt.Printf("❌ [%s] 创建临时目录失败: %v\n", clientAddr, err)
		return
	}

	// 📄 创建分片临时文件
	tempFileName := filepath.Join(tempDir, fmt.Sprintf("chunk_%d.dat", chunkIndex))
	tempFile, err := os.Create(tempFileName)
	if err != nil {
		fmt.Printf("❌ [%s] 创建临时文件失败: %v\n", clientAddr, err)
		return
	}
	defer tempFile.Close()

	// 🚀 零拷贝核心：直接从网络连接拷贝到文件，无需中间缓存
	written, err := io.CopyN(tempFile, reader, int64(chunkSize))
	if err != nil {
		fmt.Printf("❌ [%s] 读取分片数据失败: %v\n", clientAddr, err)
		return
	}
	
	if written != int64(chunkSize) {
		fmt.Printf("❌ [%s] 分片数据不完整: 期望%d, 实际%d\n", clientAddr, chunkSize, written)
		return
	}

	// 🔒 保存分片信息并增加计数
	task.mu.Lock()
	task.ReceivedChunks[chunkIndex] = tempFile
	task.ReceivedCount++
	receivedCount := task.ReceivedCount
	
	// 📊 计算接收进度并更新进度条
	progress := (receivedCount * 100) / task.TotalChunks
	if progress != task.LastProgress && (progress%5 == 0 || progress == 100 || receivedCount == 1) {
		task.LastProgress = progress
		fmt.Printf("📥 [%s] 文件 %s 接收进度: %d/%d (%d%%)\n", 
			clientAddr, task.FileName, receivedCount, task.TotalChunks, progress)
	}
	
	task.mu.Unlock()

	// ✅ 发送接收确认
	_, err = conn.Write([]byte("CHUNK_RECEIVED\n"))
	if err != nil {
		fmt.Printf("❌ [%s] 发送接收确认失败: %v\n", clientAddr, err)
		return
	}
}

// 📢 处理传输完成消息
func handleTransferComplete(conn net.Conn, clientAddr string, parts []string) {
	if len(parts) != 2 {
		fmt.Printf("❌ [%s] 传输完成消息格式错误\n", clientAddr)
		return
	}

	// 🔢 解析任务ID
	taskID := parts[1]

	// 🔍 查找对应的传输任务
	tasksMutex.RLock()
	task, exists := transferTasks[taskID]
	tasksMutex.RUnlock()
	
	if !exists {
		fmt.Printf("❌ [%s] 未找到对应的传输任务\n", clientAddr)
		return
	}

	fmt.Printf("🎉 [%s] 接收到来自客户端的传输完成通知: %s\n", clientAddr, task.FileName)
	
	// 启动合并过程
	tempDir := filepath.Join(TEMP_DIR, fmt.Sprintf("%s_%s", task.FileName, taskID))
	go mergeFile(task, clientAddr, tempDir, taskID)
	
	// ✅ 发送确认
	_, err := conn.Write([]byte("COMPLETION_ACK\n"))
	if err != nil {
		fmt.Printf("❌ [%s] 发送完成确认失败: %v\n", clientAddr, err)
		return
	}
}

// 🔗 合并文件
func mergeFile(task *TransferTask, clientAddr, tempDir string, taskID string) {
	startTime := time.Now()
	
	// 📄 创建目标文件（保存在当前目录，使用原始文件名）
	destPath := task.FileName
	destFile, err := os.Create(destPath)
	if err != nil {
		fmt.Printf("❌ [%s] 创建目标文件失败: %v\n", clientAddr, err)
		// 失败时清理临时目录
		os.RemoveAll(tempDir)
		return
	}
	defer destFile.Close()

	fmt.Printf("🔧 [%s] 开始流式合并文件: %s\n", clientAddr, task.FileName)
	
	// 🔀 按顺序合并所有分片
	buffer := make([]byte, BUFFER_SIZE)
	for i := 0; i < task.TotalChunks; i++ {
		// 📂 打开分片文件
		chunkPath := filepath.Join(tempDir, fmt.Sprintf("chunk_%d.dat", i))
		chunkFile, err := os.Open(chunkPath)
		if err != nil {
			fmt.Printf("❌ [%s] 打开分片文件失败: %v\n", clientAddr, err)
			// 失败时清理临时目录
			os.RemoveAll(tempDir)
			return
		}
		
		// 🔄 流式拷贝分片内容
		_, err = io.CopyBuffer(destFile, chunkFile, buffer)
		chunkFile.Close()
		if err != nil {
			fmt.Printf("❌ [%s] 合并分片%d失败: %v\n", clientAddr, i, err)
			// 失败时清理临时目录
			os.RemoveAll(tempDir)
			return
		}
		
		// 🧹 删除已合并的分片文件
		os.Remove(chunkPath)
	}

	// 🧹 删除临时目录（无论成功或失败都应该删除）
	defer os.RemoveAll(tempDir)
	
	fmt.Printf("🗑️  [%s] 临时目录已清理\n", clientAddr)
	
	// 📊 计算合并耗时
	elapsed := time.Since(startTime)
	fmt.Printf("✅ [%s] 流式合并完成，用时: %.2fs\n", clientAddr, elapsed.Seconds())

	// 🔐 MD5校验
	fmt.Printf("🔐 [%s] 开始MD5校验: %s\n", clientAddr, task.FileName)
	calculatedMD5, err := calculateFileMD5(destPath)
	if err != nil {
		fmt.Printf("❌ [%s] 计算MD5失败: %v\n", clientAddr, err)
		return
	}
	
	fmt.Printf("🔍 [%s] 期望MD5: %s\n", clientAddr, task.MD5Hash)
	fmt.Printf("🔍 [%s] 实际MD5: %s\n", clientAddr, calculatedMD5)

	// ✅ 校验结果
	if calculatedMD5 == task.MD5Hash {
		fmt.Printf("🎉 [%s] 文件传输成功: %s\n", clientAddr, task.FileName)
	} else {
		fmt.Printf("❌ [%s] MD5校验失败，删除文件: %s\n", clientAddr, task.FileName)
		os.Remove(destPath)
	}

	// 🧹 清理传输任务
	tasksMutex.Lock()
	delete(transferTasks, taskID)
	tasksMutex.Unlock()
}

// 🔢 计算文件MD5
func calculateFileMD5(filePath string) (string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	hash := md5.New()
	if _, err := io.Copy(hash, file); err != nil {
		return "", err
	}

	return fmt.Sprintf("%x", hash.Sum(nil)), nil
}

// 🎲 生成随机任务ID
func generateTaskID() string {
	rand.Seed(time.Now().UnixNano())
	return fmt.Sprintf("%d%d", time.Now().UnixNano(), rand.Intn(10000))
}